#include "numuCCPXSelection.hxx"
#include "numuCCSelection.hxx"
#include "baseSelection.hxx"
#include "CutUtils.hxx"


numuCCPXSelection::numuCCPXSelection(bool forceBreak): SelectionBase(forceBreak) 
{
  // define the steps
  DefineSteps();  
}

void numuCCPXSelection::DefineSteps()
{
    // Cuts must be added in the right order
    // last "true" means the step sequence is broken if cut is not passed (default is "false")
    // Add inclusive CC steps to all branches

    AddStep(StepBase::kCut,    "event quality",        new EventQualityCut(),           true);
    AddStep(StepBase::kCut,    "> 0 tracks ",          new TotalMultiplicityCut(),      true);  
    AddStep(StepBase::kAction, "find leading tracks",  new FindLeadingTracksAction());  

    AddStep(StepBase::kAction, "find tracker tracks",  new FindTrackerTracksAction()); //local action; on std_numuccPX_actions::FindTracksWithTPC, std_numuccPX_actions::FindTracksWithFGDAndNoTPC
    AddStep(StepBase::kAction, "find candidates",      new FindCandidatesAction()); //local action;

    // First branch is for cases 0,1,2
    AddStep(StepBase::kAction, "find vertex",        new FindVertexAction());  
    AddStep(StepBase::kCut,    "quality+fiducial",   new TrackQualityFiducialCut(),   true);  
    AddStep(StepBase::kAction, "find veto track",    new FindVetoTrackAction());
    AddStep(StepBase::kCut,    "veto",               new ExternalVetoCut());
    AddStep(StepBase::kAction, "find oofv track",    new FindOOFVTrackAction());
    AddStep(StepBase::kCut,    "muon PID",           new MuonPIDCut());

    AddStep(StepBase::kCut,    "1 neg track",        new OneNegativeTrackCut()); //local
    AddStep(StepBase::kAction, "find michel e-",     new FindMichelElectronsAction()); //local
    AddStep(StepBase::kCut,    "no michel ele",      new MichelElectronCut()); //local
    AddStep(StepBase::kCut,    "two tracker tracks",       new TwoTrackerTracksCut());//local
    AddStep(StepBase::kAction, "find correct sense",       new FindCorrectProtonSenseAction()); //local

    AddStep(StepBase::kCut, "1 pos track",        new OnePositiveTrackCut()); //local
    AddStep(StepBase::kCut, "common vertex",      new CommonVertexCut1()); //local
    AddStep(StepBase::kCut, "proton PID",         new ProtonPIDCut()); //local

    //very important! without it no events shown in Scan
    SetBranchAlias(0,"muTPC-pTPC");
}


bool numuCCPXSelection::FillEventSummary(const AnaBoxB& box, AnaEventB& event, Int_t allCutsPassed[])
{
    (void)box;

    // Should set event sample enumeration here
    if(allCutsPassed[0]) event.Summary->EventSample = gkSampleEnumCCPX;
    return (event.Summary->EventSample != nd280Samples::kUnassigned);
}


bool OneNegativeTrackCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    // Exactly one negative track with good quality and in FV
    if (box.nNegativeTPCtracks==1) 
        return true;
    else 
        return false;
}

bool MichelElectronCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    // Temporarily, a Float_t cut to make sure it works with OldFlatTree, FlatTree and oaAnalysis
    if (box.nFGD1MichelElectrons == 0)
        return true;
    else
        return false;
}

bool TwoTrackerTracksCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    if (ccPXbox->TrackerTracks.size() == 2) 
        return true;
    else 
        return false;
}


bool OnePositiveTrackCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    // Exactly one positive track with good quality and in FV
    if (box.nPositiveTPCtracks == 1) 
        return true;
    else 
        return false;
}

bool CommonVertexCut1::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    if (!ccPXbox->MuonCandidate || !ccPXbox->ProtonCandidate) return false;

    return cutUtils::CommonVertexCut(*(ccPXbox->MuonCandidate), *(ccPXbox->ProtonCandidate), 
            _cutX, _cutY, _cutZ);
}


bool ProtonPIDCut::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    if (!box.HMPtrack) return false;

    // use different cut value w.r.t. to cutUtils::ProtonPIDCut()  
    if (anaUtils::GetPIDLikelihood(*(box.HMPtrack),2)>0.5)
        return true;
    else
        return false;
}

bool FindMichelElectronsAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    cutUtils::FindFGDMichelElectrons(event,box);

    return true;
}

bool FindCandidatesAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    if (box.HMNtrack)
        ccPXbox->MuonCandidate = box.HMNtrack;
    else if (ccPXbox->FGDtracksCrossingFGD1.size()>0)
        ccPXbox->MuonCandidate = ccPXbox->FGDtracksCrossingFGD1[0]; //tracks crossing FGD1

    if (box.HMPtrack)
        ccPXbox->ProtonCandidate = box.HMPtrack;
    else if (box.HMNtrack && ccPXbox->FGDtracksCrossingFGD1.size()>0)
        ccPXbox->ProtonCandidate = ccPXbox->FGDtracksCrossingFGD1[0];  //tracks crossing FGD1

    return true;
}

bool FindTrackerTracksAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    // all Tracker tracks
    ccPXbox->TrackerTracks.clear();

    std_numuccPX_actions::FindTracksWithTPC(event,*ccPXbox, true, SubDetId::kFGD1);
    //needed!! otherwise the selected sample is different!! MAYBE due to michel electron cut
    std_numuccPX_actions::FindTracksWithFGDAndNoTPC(event, *ccPXbox, SubDetId::kFGD1);

    return true;
}

bool FindCorrectProtonSenseAction::Apply(const AnaEventB& event, AnaBoxB& box) const
{
    (void)event;

    AnaBoxCCPX * ccPXbox = static_cast<AnaBoxCCPX*>(&box); 

    if (ccPXbox->ProtonCandidate)
        // it must reverse the track only if it is a FGD track (no TPC segments)
        if (!anaUtils::TrackUsesDet(*(ccPXbox->ProtonCandidate), SubDetId::kTPC))
            std_numuccPX_actions::FindCorrectSense(ccPXbox->ProtonCandidate, box.Vertex);

    return true;
}

void std_numuccPX_actions::FindTracksWithTPC(const AnaEventB& event, AnaBoxCCPX& ccPXbox, bool qualityCut, SubDetId::SubDetEnum det)
{
  std::vector<AnaTrackB*> allTracks[5];

  // all tracks using TPC
  AnaTrackB* selTracks[100];
  int nTPC = anaUtils::GetAllTracksUsingTPC(event, selTracks);

  //loop over tpc tracks
  for (Int_t i=0;i<nTPC; ++i){
    AnaTrackB* track = selTracks[i];

    // all TPC tracks
    allTracks[4].push_back(track);

    // all Tracker tracks
    ccPXbox.TrackerTracks.push_back(track);

    //check if the track starts in FGD1 volume
    if (anaUtils::InDetVolume(SubDetId::kFGD1,track->PositionStart)) {
      if (qualityCut){
        //apply the quality cut: good quality tracks
        if (cutUtils::TrackQualityCut(*track)) { 
          // Apply the fiducial cut: save tracks out FV
          if (!cutUtils::FiducialCut(*track, det)) {

            if (track->Charge == -1)
              allTracks[0].push_back(track);
            else
              allTracks[1].push_back(track);
          }
        }
        //bad quality tracks
        else {
          // Apply the fiducial cut: save tracks in FV
          if (cutUtils::FiducialCut(*track, det))
            allTracks[2].push_back(track);
          // save tracks out FV
          else
            allTracks[3].push_back(track);
        }
      }
    }
  }

  // Fill the box
  ccPXbox.NegativeTPCtracksOutFV = allTracks[0];
  ccPXbox.PositiveTPCtracksOutFV = allTracks[1];
  ccPXbox.BadQuaTPCtracksInFV = allTracks[2];
  ccPXbox.BadQuaTPCtracksOutFV = allTracks[3];
  ccPXbox.TPCtracks = allTracks[4];
}

void std_numuccPX_actions::FindTracksWithFGDAndNoTPC(const AnaEventB& event, AnaBoxCCPX& ccPXbox, SubDetId::SubDetEnum det)
{
  std::vector<AnaTrackB*> allTracks[7];

  // Get all tracks with FGD and no TPC
  AnaTrackB* allFgdTracks[100];
  int nFgdTracks= anaUtils::GetAllTracksUsingFGDAndNoTPC(event, allFgdTracks);

  //loop over tracks
  for (int i=0;i<nFgdTracks;i++){
    AnaTrackB* track = allFgdTracks[i];

    // All tracks using FGD and no TPC
    allTracks[6].push_back(track);      

    // All Tracker tracks
    ccPXbox.TrackerTracks.push_back(track);

    // Tracks using FGD1 (and no TPC)
    if (anaUtils::TrackUsesDet(*track,SubDetId::kFGD1)) 
      allTracks[5].push_back(track);      

    //check if the track starts in FGD1 volume
    if (anaUtils::InDetVolume(SubDetId::kFGD1,track->PositionStart)) {
      allTracks[0].push_back(track);      

      // in FV tracks
      if (cutUtils::FiducialCut(*track, det)) {
        // contained FGD tracks
        if (anaUtils::TrackUsesOnlyDet(*track, SubDetId::kFGD1))
          allTracks[1].push_back(track);
        // not-contained FGD tracks
        else
          allTracks[2].push_back(track);      
      }
      // out FV tracks
      else {
        // contained FGD tracks
        if (anaUtils::TrackUsesOnlyDet(*track, SubDetId::kFGD1))
          allTracks[3].push_back(track);
        // not-contained FGD tracks
        else 
          allTracks[4].push_back(track);      
      }
    }
  }

  ccPXbox.FGDtracksInFGD1Volume = allTracks[0];
  ccPXbox.FGDContTracksInFV = allTracks[1];
  ccPXbox.FGDNotContTracksInFV = allTracks[2];
  ccPXbox.FGDContTracksOutFV = allTracks[3];
  ccPXbox.FGDNotContTracksOutFV = allTracks[4];
  ccPXbox.FGDtracksCrossingFGD1 = allTracks[5];
  ccPXbox.FGDtracks = allTracks[6];
}

void std_numuccPX_actions::FindCorrectSense(AnaTrackB* track, AnaVertexB* vertex) 
{
  if (!track) return;
  if (!vertex) return;

  Float_t distStart = cutUtils::GetSeparationSquared(track->PositionStart, vertex->Position);
  Float_t distEnd   = cutUtils::GetSeparationSquared(track->PositionEnd, vertex->Position);

  if (distEnd < distStart){
    // reverse the track
    Float_t posTemp[4] = {track->PositionStart[0], track->PositionStart[1], track->PositionStart[2], track->PositionStart[3]};
    //    TVector3       dirTemp = track->DirectionStart;
    cutUtils::AssignArray(track->PositionEnd, track->PositionStart, 4);
    cutUtils::AssignArray(posTemp, track->PositionEnd, 4);
    //    track->DirectionStart   = - track->DirectionEnd;
    //    track->DirectionEnd     = - dirTemp;    
    track->Charge           *= -1;    
  }

}
